/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include "OriginalNamePass.h"

#include <vector>

#include "ClassHierarchy.h"
#include "ConfigFiles.h"
#include "DexAnnotation.h"
#include "DexLimitsInfo.h"

#include "DexStoreUtil.h"

#include "DexUtil.h"
#include "PassManager.h"
#include "Show.h"
#include "Trace.h"

#define METRIC_MISSING_ORIGINAL_NAME_ROOT "num_missing_original_name_root"
#define METRIC_ORIGINAL_NAME_COUNT "num_original_name"

static OriginalNamePass s_pass;

static const char* redex_field_name = "__redex_internal_original_name";

void OriginalNamePass::build_hierarchies(
    PassManager& mgr,
    const ClassHierarchy& ch,
    Scope& scope,
    UnorderedMap<const DexType*, std::string_view>* hierarchies) {
  std::vector<DexClass*> base_classes;
  for (const auto& base : m_hierarchy_roots) {
    // skip comments
    if (base.c_str()[0] == '#') {
      continue;
    }
    auto* base_type = DexType::get_type(base);
    auto* base_class = base_type != nullptr ? type_class(base_type) : nullptr;
    if (base_class == nullptr) {
      TRACE(ORIGINALNAME, 2,
            "Can't find class for annotate_original_name rule %s",
            base.c_str());
      mgr.incr_metric(METRIC_MISSING_ORIGINAL_NAME_ROOT, 1);
    } else {
      base_classes.emplace_back(base_class);
    }
  }
  for (const auto& base_class : base_classes) {
    auto base_name = base_class->get_deobfuscated_name_or_empty();
    hierarchies->emplace(base_class->get_type(), base_name);
    TypeSet children_and_implementors;
    get_all_children_or_implementors(ch, scope, base_class,
                                     children_and_implementors);
    for (const auto& cls : children_and_implementors) {
      hierarchies->emplace(cls, base_name);
    }
  }
}

void OriginalNamePass::run_pass(DexStoresVector& stores,
                                ConfigFiles& conf,
                                PassManager& mgr) {
  auto scope = build_class_scope(stores);
  ClassHierarchy ch = build_type_hierarchy(scope);
  UnorderedMap<const DexType*, std::string_view> to_annotate;
  build_hierarchies(mgr, ch, scope, &to_annotate);
  const auto* field_name = DexString::make_string(redex_field_name);
  DexType* string_type = type::java_lang_String();
  init_classes::InitClassesWithSideEffects init_classes_with_side_effects(
      scope, conf.create_init_class_insns());

  size_t store_id = 0;
  for (auto& store : stores) {
    // Backup dex, in case we exceed field limits. We assuem there is at most 1
    // new dex for each store.
    DexClasses new_dex;
    bool canary_inserted = false;
    size_t dex_id = 0;
    for (auto& dex : store.get_dexen()) {
      DexLimitsInfo dex_limits(&init_classes_with_side_effects);
      DexClasses overflow_classes;
      for (auto const& cls : dex) {
        DexType* cls_type = cls->get_type();
        if ((to_annotate.count(cls_type) == 0u) ||
            !cls->rstate.is_renamable_initialized_and_renamable() ||
            cls->rstate.is_generated()) {
          // No need to keep original name.
          if (!dex_limits.update_refs_by_adding_class(cls)) {
            // Move cls from current dex to new_dex.
            overflow_classes.emplace_back(cls);
          }
          continue;
        }

        always_assert_log(!cls->rstate.is_generated(),
                          "Encountered a generated class %s in %s. The class "
                          "may be generated by "
                          "merging other targeting classes by ClassMerging "
                          "service, disable the "
                          "merging to fix the error.",
                          SHOW(cls->get_deobfuscated_name_or_empty()),
                          SHOW(this->name()));
        auto external_name = java_names::internal_to_external(
            cls->get_deobfuscated_name_or_empty());
        auto lastDot = external_name.find_last_of('.');
        auto simple_name = (lastDot != std::string::npos)
                               ? external_name.substr(lastDot + 1)
                               : external_name;
        const auto* simple_name_s = DexString::make_string(simple_name);
        always_assert_log(
            DexField::get_field(cls_type, field_name, string_type) == nullptr,
            "field %s already exists!",
            redex_field_name);

        always_assert_log(
            DexField::get_field(cls_type, field_name, string_type) == nullptr,
            "field %s already exists!",
            redex_field_name);
        DexField* f =
            DexField::make_field(cls_type, field_name, string_type)
                ->make_concrete(ACC_PUBLIC | ACC_STATIC | ACC_FINAL,
                                std::unique_ptr<DexEncodedValue>(
                                    new DexEncodedValueString(simple_name_s)));
        // These fields are accessed reflectively, so make sure we do not remove
        // them.
        f->rstate.set_root();
        insert_sorted(cls->get_sfields(), f, compare_dexfields);
        f->set_deobfuscated_name(show_deobfuscated(f));

        if (!dex_limits.update_refs_by_adding_class(cls)) {
          overflow_classes.emplace_back(cls);
        }

        mgr.incr_metric(METRIC_ORIGINAL_NAME_COUNT, 1);
        mgr.incr_metric(std::string(METRIC_ORIGINAL_NAME_COUNT) +
                            "::" + to_annotate[cls_type],
                        1);
      }
      TRACE(ORIGINALNAME, 2,
            "Now Dex %zu in store %zu has been done and orignial size is %zu\n",
            dex_id, store_id, dex.size());
      if (!overflow_classes.empty()) {
        // Remove those classes from current dex.
        size_t removed = std::erase_if(dex, [&](auto* cls) {
          return std::find(overflow_classes.begin(), overflow_classes.end(),
                           cls) != overflow_classes.end();
        });
        TRACE(
            ORIGINALNAME, 2,
            "%zu (%zu)  classes has been removed from %zu of store %s, and the "
            "rest number is %zu\n",
            overflow_classes.size(), removed, dex_id, store.get_name().c_str(),
            dex.size());
        if (new_dex.empty()) {
          // A canary_cls need to be added when a new dex is created.
          int dexnum = static_cast<int>(store.get_dexen().size());
          auto store_name = store.is_root_store() ? "" : store.get_name();
          DexClass* canary_cls = create_canary(dexnum, store_name);
          for (auto* m : canary_cls->get_all_methods()) {
            if (m->get_code() == nullptr) {
              continue;
            }
            m->get_code()->build_cfg();
          }
          always_assert(!canary_inserted);
          new_dex.push_back(canary_cls);
          canary_inserted = true;
        }
        new_dex.insert(new_dex.end(), overflow_classes.begin(),
                       overflow_classes.end());
      }
      dex_id++;
    }
    if (!new_dex.empty()) {
      // A new dex is created due to too many field refs.
      TRACE(ORIGINALNAME, 2,
            "Before adding a new dex, current %zu dexes in "
            "store %s\n",
            store.num_dexes(), store.get_name().c_str());
      store.add_classes(new_dex);
      TRACE(ORIGINALNAME, 2,
            "%zu classes has been added into a new dex, current %zu dexes in "
            "store %s\n",
            new_dex.size(), store.num_dexes(), store.get_name().c_str());
    }
  }
}
